package com.ibm.iisp.common.datasource;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

@Data
@Slf4j
public class DynamicDataSource extends AbstractDataSource {
    private DataSource master;
    private List<DataSource> slaves = new ArrayList<>();
    private AtomicInteger pointer = new AtomicInteger(-1);
    protected DataSource determineTargetDataSource() {
        boolean readonly = TransactionSynchronizationManager.isCurrentTransactionReadOnly();
        log.debug("isCurrentTransactionReadOnly: {}", readonly);
        if (readonly) {
            if (slaves.size() > 1) {
                int i = pointer.incrementAndGet();
                if (i > Integer.MAX_VALUE / 2) {
                    pointer.set(0);
                    i = 0;
                }
                int index = i % slaves.size();
                log.trace("return readonly ds: {}", index);
                return slaves.get(index);
            } else if (slaves.size() == 1) {
                return slaves.get(0);
            }
        }
        return master;

    }

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }
}
